"""
Navier-Stokes Matrix Assembly

Assembles system matrices for the Navier-Stokes equations with multi-domain
coupling and spectral element discretization.
"""

"""
    generate_system_matrices(multidomain::MultiDomain, D2::AbstractMatrix, verbose::Bool=false)

Generate system matrices for Navier-Stokes equations.
"""
function generate_system_matrices(multidomain::MultiDomain, D2::AbstractMatrix, verbose::Bool=false)
    
    if verbose
        println("Generating system matrices...")
        println("  Multi-domain blocks: $(multidomain.index_sum)")
        println("  Polynomial order: $(multidomain.n)")
    end
    
    # Generate identity and boundary matrices
    I, I_full, I_inner, I_bound, Ib = generate_identity_matrices(multidomain.n)
    
    # Create boundary condition matrices
    Dn0, Dn1, Dn2, Dn3 = create_boundary_matrices(I_inner)
    
    # Construct projection matrices
    Proj = construct_projection_matrices(multidomain.matrix_dist, Dn0, Dn1, Dn2, Dn3, multidomain.g)
    
    # Create matrices for velocity and pressure equations
    I_fix_u = multidomain.dim == 3 ? I_full / 3 : I_full / 2
    Dm_u = I_fix_u - multidomain.time_param * D2
    Dm = copy(D2)  # For pressure equation
    
    if verbose
        println("  Velocity matrix size: $(size(Dm_u))")
        println("  Pressure matrix size: $(size(Dm))")
        println("  Projection matrices: $(length(Proj)) components")
    end
    
    return Dm, Dm_u, Proj
end

"""
    generate_identity_matrices(n::Int)

Generate identity and boundary-related matrices for spectral elements.
"""
function generate_identity_matrices(n::Int)
    N = n
    I = ones(Float64, N+1)  # 1D identity vector
    I_full = Matrix{Float64}(LinearAlgebra.I, N+1, N+1)  # Full identity matrix
    
    # Interior matrix (excluding boundaries)
    I_inner = copy(I_full)
    I_inner[1, 1] = 0.0
    I_inner[end, end] = 0.0
    
    # Boundary matrix
    I_bound = I_full - I_inner
    
    # Boundary vector
    Ib = ones(Float64, (N-1))  # For interior points only
    
    return I, I_full, I_inner, I_bound, Ib
end

"""
    create_boundary_matrices(I_inner::Matrix{Float64})

Create boundary condition matrices for different boundary types.
"""
function create_boundary_matrices(I_inner::Matrix{Float64})
    # Create copies for different boundary conditions
    Dn0 = copy(I_inner)  # Left boundary
    Dn1 = copy(I_inner)  # Right boundary  
    Dn2 = copy(I_inner)  # Bottom boundary
    Dn3 = copy(I_inner)  # Top boundary
    
    # Modify boundary conditions
    Dn0[1, 1] = 1.0
    Dn1[1, 1] = 1.0
    Dn2[end, end] = 1.0
    Dn3[end, end] = 1.0
    
    return Dn0, Dn1, Dn2, Dn3
end

"""
    construct_projection_matrices(matrix_dist::Array{Int}, Dn0, Dn1, Dn2, Dn3, g)

Construct projection matrices for multi-domain coupling.
"""
function construct_projection_matrices(matrix_dist::Array{Int}, Dn0, Dn1, Dn2, Dn3, g)
    dim = length(size(matrix_dist))
    index_sum = length(g)
    n = size(Dn0, 1) - 1
    
    # Initialize projection matrices
    if dim == 2
        P0 = zeros(Float64, index_sum + 1, n + 1, n + 1)
        P1 = zeros(Float64, index_sum + 1, n + 1, n + 1)
        Proj = (P0, P1)
    else  # 3D
        P0 = zeros(Float64, index_sum + 1, n + 1, n + 1)
        P1 = zeros(Float64, index_sum + 1, n + 1, n + 1)
        P2 = zeros(Float64, index_sum + 1, n + 1, n + 1)
        Proj = (P0, P1, P2)
    end
    
    # Dictionary of boundary matrices
    dn_matrices = Dict(0 => Dn0, 1 => Dn1, 2 => Dn2, 3 => Dn3)
    
    # Process each active block
    for l in 1:index_sum
        if dim == 2
            i0, i1 = g[l][1], g[l][2]
            
            # X direction projection
            t = matrix_dist[i0-1, i1] + 2 * matrix_dist[i0+1, i1]
            P0[l+1, :, :] = dn_matrices[Int(t)]
            
            # Y direction projection
            t = matrix_dist[i0, i1-1] + 2 * matrix_dist[i0, i1+1]
            P1[l+1, :, :] = dn_matrices[Int(t)]
            
        else  # 3D
            i0, i1, i2 = g[l][1], g[l][2], g[l][3]
            
            # X direction projection
            t = matrix_dist[i0-1, i1, i2] + 2 * matrix_dist[i0+1, i1, i2]
            P0[l+1, :, :] = dn_matrices[Int(t)]
            
            # Y direction projection
            t = matrix_dist[i0, i1-1, i2] + 2 * matrix_dist[i0, i1+1, i2]
            P1[l+1, :, :] = dn_matrices[Int(t)]
            
            # Z direction projection
            t = matrix_dist[i0, i1, i2-1] + 2 * matrix_dist[i0, i1, i2+1]
            P2[l+1, :, :] = dn_matrices[Int(t)]
        end
    end
    
    return Proj
end

"""
    generate_Lsp_matrix(Dm::AbstractMatrix, Proj::Tuple, matrix_dist::Array{Int}, 
                       g, gi, multidomain::MultiDomain, title::String="Lsp")

Generate the linear system operator matrix for pressure or velocity equations.
"""
function generate_Lsp_matrix(Dm::AbstractMatrix, Proj::Tuple, matrix_dist::Array{Int}, 
                            g, gi, multidomain::MultiDomain, title::String="Lsp")
    
    N = multidomain.n
    NP2 = multidomain.n_plus_1^multidomain.dim
    index_sum = multidomain.index_sum
    
    # Initialize block matrices
    block_diag_matrix = zeros(Float64, NP2 * index_sum, NP2 * index_sum)
    block_diag_matrixT = similar(block_diag_matrix)
    
    # Generate identity matrices for tensor operations
    I, I_full, I_inner, I_bound, Ib = generate_identity_matrices(N)
    
    # Create dimension-specific differential operator
    if multidomain.dim == 3
        Dx = kron(kron(Dm, I_full), I_full) + kron(kron(I_full, Dm), I_full) + kron(kron(I_full, I_full), Dm)
    else  # 2D
        Dx = kron(Dm, I_full) + kron(I_full, Dm)
    end
    
    # Create modified Dm with boundary rows zeroed out
    Dmc = copy(Dm)
    Dmc[1, :] .= 0.0
    Dmc[end, :] .= 0.0
    
    # Process each active block
    for l in 1:index_sum
        gn = l + 1  # Block index (1-based)
        
        # Get grid indices
        if multidomain.dim == 3
            i0, i1, i2 = g[l][1], g[l][2], g[l][3]
            # 3D implementation would go here
        else  # 2D
            i0, i1 = g[l][1], g[l][2]
            
            # Create main operator for current block
            coo = (kron(Dmc, I_inner) + kron(I_inner, Dmc) + 
                   Matrix{Float64}(LinearAlgebra.I, NP2, NP2) - 
                   kron(Proj[1][gn, :, :], Proj[2][gn, :, :]))
            
            # Add to block diagonal matrix T  
            seq_vec = create_sequence_vector(gn, index_sum)
            block_diag_matrixT += kron(kron(seq_vec, seq_vec'), coo)
        end
    end
    
    # Combine matrices
    Lsp = block_diag_matrixT + block_diag_matrix
    
    # Add regularization for numerical stability
    reg = 1e-3 * Matrix{Float64}(LinearAlgebra.I, size(Lsp, 1), size(Lsp, 2))
    Lsp += reg
    
    # Print eigenvalue information for diagnostics
    if size(Lsp, 1) <= 1000  # Only for small matrices
        try
            eigenvalues = eigvals(Lsp' * Lsp)
            sorted_eigs = sort(real.(eigenvalues))
            println("$title eigenvalues: smallest=$(sorted_eigs[1:2]), largest=$(sorted_eigs[end-1:end])")
        catch
            println("Could not compute eigenvalues for $title matrix")
        end
    end
    
    return Lsp
end

"""
    create_sequence_vector(gn::Int, index_sum::Int) -> Vector{Float64}

Create a one-hot sequence vector for Kronecker products.
"""
function create_sequence_vector(gn::Int, index_sum::Int)
    sequence = zeros(Float64, index_sum, 1)
    if 1 <= gn <= index_sum
        sequence[gn] = 1.0
    end
    return sequence
end

"""
    assemble_ns_matrices(multidomain::MultiDomain{D}, sem_operators::SEMOperators{D},
                        backend::Symbol=:julia) where D

Assemble complete Navier-Stokes system matrices.
"""
function assemble_ns_matrices(multidomain::MultiDomain{D}, sem_operators::SEMOperators{D},
                             backend::Symbol=:julia) where D
    
    # Extract derivative matrices
    D1, D2 = sem_operators.D1, sem_operators.D2
    
    # Generate system matrices
    Dm, Dm_u, Proj = generate_system_matrices(multidomain, D2, true)
    
    # Generate linear system operators
    Lsp_u = generate_Lsp_matrix(Dm_u, Proj, multidomain.matrix_dist, 
                               multidomain.g, multidomain.gi, multidomain, "Lsp_u")
    Lsp = generate_Lsp_matrix(Dm, Proj, multidomain.matrix_dist, 
                             multidomain.g, multidomain.gi, multidomain, "Lsp")
    
    return (
        velocity_matrix = Lsp_u,
        pressure_matrix = Lsp, 
        derivative_matrices = (D1=D1, D2=D2),
        projection_matrices = Proj,
        multidomain = multidomain
    )
end